import rasterio
import numpy as np
import math
from tqdm import tqdm

# todo 电脑可以处理，暂时没必要分块：

def get_LS(sp):
    if sp <= 1.0:
        m = 2
    elif 1.0 < sp <= 3.0:
        m = 0.3
    elif 3.0 < sp <= 5.0:
        m = 0.4
    else:
        m = 0.5

    length = 30.0 / math.cos(sp / 180 * math.pi)
    lv = (length / 22.13) ** m

    if sp < 5.0:
        sv = 10.8 * math.sin(sp / 180 * math.pi) + 0.03
    elif 5.0 <= sp <= 10.0:
        sv = 16.8 * math.sin(sp / 180 * math.pi) - 0.5
    else:
        sv = 21.9 * math.sin(sp / 180 * math.pi) - 0.96

    return lv, sv


def calculate_lv_sv(input_file_path, output_lv_path, output_sv_path, block_size=512):
    with rasterio.open(input_file_path) as src:
        # 获取图像元数据
        meta = src.meta.copy()
        meta.update(dtype=rasterio.float32)

        # 设置分块大小
        meta['blockxsize'] = block_size
        meta['blockysize'] = block_size

        # 创建输出文件
        with rasterio.open(output_lv_path, 'w', **meta) as dst_lv, \
             rasterio.open(output_sv_path, 'w', **meta) as dst_sv:

            windows = list(src.block_windows())
            total_blocks = len(windows)

            for ji, window in tqdm(windows, desc="Processing blocks", unit="block"):
                slop_data = src.read(1, window=window)

                # 初始化输出数组
                lv_data = np.zeros(slop_data.shape).astype(np.float32)
                sv_data = np.zeros(slop_data.shape).astype(np.float32)

                # 计算lv和sv
                for i in range(slop_data.shape[0]):
                    for j in range(slop_data.shape[1]):
                        if not np.isnan(slop_data[i, j]):
                            lv, sv = get_LS(slop_data[i, j])
                            lv_data[i, j] = lv
                            sv_data[i, j] = sv

                # 写入结果
                dst_lv.write(lv_data.astype(rasterio.float32), 1, window=window)
                dst_sv.write(sv_data.astype(rasterio.float32), 1, window=window)


# 示例用法
input_file_path = r'H:\GEP_data\slop_merged.tif'  # 替换为您的坡度图像文件路径
output_lv_path = r'H:\GEP_data\插值好的数据\SR\SR_L_china_30m.tif'
output_sv_path = r'H:\GEP_data\插值好的数据\SR\SR_S_china_30m.tif'
calculate_lv_sv(input_file_path, output_lv_path, output_sv_path, block_size=102400)
"""
block_size

假设有以下条件：
数据类型为 32位浮点数（float32），占用 4 字节。
图像是单波段的。
block_size 为 1024。
那么每个块的数据量计算如下：
每个块的数据量=1024×1024×4 字节
每个块的数据量=4,194,304 字节
每个块的数据量=4 MB

block_size=10240000 40G
"""